import matplotlib.pyplot as plt
import torch


def _set_lim(xlim, ylim):
    if isinstance(xlim, int) or isinstance(xlim, float):
        plt.xlim([0, xlim])
    elif xlim is not None:
        plt.xlim(list(xlim))
    if isinstance(ylim, int) or isinstance(ylim, int):
        plt.ylim([0, ylim])
    elif ylim is not None:
        plt.ylim(list(ylim))


@torch.no_grad()
def draw(train_loss, train_accuracy, test_loss, test_accuracy, xlim=None, ylim_loss=None, ylim_accuracy=None):
    plt.figure(figsize=(13, 4.8)).suptitle(f'Result')
    plt.subplot(121)
    _set_lim(xlim, ylim_loss)
    plt.plot(train_loss, color='blue', label='train loss')
    plt.plot(test_loss, color='red', label='test loss')
    plt.legend()
    plt.subplot(122)
    _set_lim(xlim, ylim_accuracy)
    plt.plot(train_accuracy, color='blue', label='train accuracy')
    plt.plot(test_accuracy, color='red', label='test accuracy')
    plt.legend()
    plt.show()
